import torch 
from torch import nn
from torch.nn import functional as F
import pickle
from typing import NamedTuple
import numpy as np
from scipy import sparse
import os
from .utils import compute_v_normals, compute_normals, batch_rodrigues

BFM_MODEL_PATH = "asset/3DMM/3DMM_info.npy"
BFM_MESH_PATH = "asset/3DMM/sub_mesh.obj"
BFM_TRI_PATH = "asset/3DMM/tris.txt"
BFM_LMK_PATH = "asset/3DMM/lands_info.txt"

# %% class
class BFMLBS(nn.Module):
  dtype = torch.float32
  
  def register_from_numpy(self, param, dtype=None):
      dtype = dtype or self.dtype
      return nn.Parameter(torch.from_numpy(param).to(dtype), requires_grad=False)
    
  def __init__(self, 
               bfm_model_path: str = BFM_MODEL_PATH,
               bfm_face_path: str = BFM_TRI_PATH,
               bfm_lands_path: str = BFM_LMK_PATH,
               id_dim:int = 100,
               exp_dim:int = 79,
               ):
    super().__init__()
    # FACE_PART dim 34650
    # include ear & neck: 53490
    
    params: dict = np.load(bfm_model_path, allow_pickle=True).item()
    lmk_info = np.loadtxt(bfm_lands_path, dtype=np.int32)
    self.lmk_info = nn.Parameter(torch.from_numpy(lmk_info), requires_grad=False)  # 68 x 3
    
    base_shape = params["b_shape"][:id_dim, :].reshape(id_dim, -1, 3)  # shape_dim x N x 3
    base_exp = params["b_exp"][:exp_dim, :].reshape(exp_dim, -1, 3)    # exp_dim x N x 3
    
    mu_shape = params["mu_shape"] / 1e5
    mu_exp = params["mu_exp"] / 1e5
    
    v_template = (mu_shape + mu_exp).reshape(-1, 3)
    v_template = v_template - v_template.mean(axis=0)
    
    self.v_template = self.register_from_numpy(v_template) # N x 3
    self.base_shape = self.register_from_numpy(base_shape) # shape_dim x N x 3
    self.base_exp = self.register_from_numpy(base_exp)     # exp_dim x N x 3
    
    # 标准差, 相当于lbs weights
    sig_shape = params["sig_shape"][:id_dim] / 1e5
    sig_exp = params["sig_exp"][:exp_dim] / 1e5
    self.sig_shape = self.register_from_numpy(sig_shape)  # shape_dim
    self.sig_exp = self.register_from_numpy(sig_exp)      # exp_dim
    
    # Faces
    tris = np.loadtxt(bfm_face_path, dtype=np.int32)
    self.faces = nn.Parameter(torch.from_numpy(tris - 1), requires_grad=False)
    
    
  def forward_sub(self, 
                    shape: torch.Tensor, 
                    expr: torch.Tensor,
                    rotation: torch.Tensor,
                    neck: torch.Tensor=None,
                    jaw: torch.Tensor=None,
                    eyes: torch.Tensor=None,
                    translation: torch.Tensor=None,
                    static_offset: torch.Tensor=None,
                    sub_index: torch.Tensor = None,
                    ):
    batch_size = expr.size(0)
    shape_para = shape * self.sig_shape
    exp_para = expr * self.sig_exp
    if sub_index is None:
      sub_index = torch.arange(self.v_template.size(0), device=shape.device).long()
    
    geo_shape = torch.einsum("bk, kji -> bji", shape_para, self.base_shape[:, sub_index])
    geo_expr = torch.einsum("bk, kji -> bji", exp_para, self.base_exp[:, sub_index])
    vertices = geo_shape + geo_expr + self.v_template[None, sub_index]
    
    rotmat = batch_rodrigues(rotation).view(batch_size, 3, 3)
    vertices_rotated = torch.einsum("bni, bji -> bnj", vertices, rotmat)
    if not translation is None: # shape should be B x 3
      vertices_rotated = vertices_rotated + translation[:, None, :]
      
    return vertices_rotated

  def forward(
            self,
            shape:torch.Tensor,
            expr:torch.Tensor,
            rotation:torch.Tensor,
            neck:torch.Tensor=None,
            jaw:torch.Tensor=None,
            eyes:torch.Tensor=None,
            translation=None,
            static_offset=None,
        ):
    batch_size = expr.size(0)
    shape_para = shape * self.sig_shape
    exp_para = expr * self.sig_exp
    
    geo_shape = torch.einsum("bk, kji -> bji", shape_para, self.base_shape)
    geo_expr = torch.einsum("bk, kji -> bji", exp_para, self.base_exp)
    vertices = geo_shape + geo_expr + self.v_template[None]
    
    rotmat = batch_rodrigues(rotation).view(batch_size, 3, 3)
    vertices_rotated = torch.einsum("bni, bji -> bnj", vertices, rotmat)
    if not translation is None: # shape should be B x 3
      vertices_rotated = vertices_rotated + translation[:, None, :]
    
    return vertices_rotated